from solutions import solver_1, solver_2, solver_3, solver_4, solver_5
import numpy as np
import os
import argparse


def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--solver', type=int, default=1,
                        help='选择不同问题的解决方案，用整数表示，1-5')
    parser.add_argument('--lr', type=float, default=1e-6, help='学习率')
    parser.add_argument('--epoch', type=int, default=40000)
    parser.add_argument('--store_path', type=str,
                        default='./result', help='存储路径')
    parser.add_argument('--is_post_train', type=bool,
                        default=True, help='是否进行后期训练，即使用遗传算法进一步逼近')
    parser.add_argument('--rank', '-t', type=int,
                        default=2, help='矩阵的阶数，这里用的是2的t次方')
    parser.add_argument('--K', type=int, default=2, help='矩阵序列的个数')
    return parser.parse_args()


if __name__ == '__main__':
    args = parse()
    solvers = [solver_1, solver_2, solver_3, solver_4, solver_5]
    M = solvers[args.solver-1](2**args.rank, args.K)
    if not os.path.exists(args.store_path):
        os.mkdir(args.store_path)
    if args.solver == 1:
        args.is_post_train = False
    M.train(lr=args.lr, epoch=args.epoch, store_path=args.store_path,
            is_post_train=args.is_post_train)
    X = np.random.random([2**args.rank, 1]).astype(complex)
    print(M.beta * M.dft - M.result)
    M.print()
